from testbase import TestBase
import rpyc
from rpyc.utils.server import ThreadedServer
from rpyc.utils.authenticators import PubkeyDBAuthenticator
from rpyc.utils.security import PubkeyDB, ProtocolAuthError
import thread, time


class PkdbTest(TestBase):
    def setup(self):
        self.log("creating server and client PKDBs")
        self.client1_pkdb = PubkeyDB.create()
        self.client2_pkdb = PubkeyDB.create()
        self.client3_pkdb = PubkeyDB.create()
        self.server_pkdb = PubkeyDB.create()
        self.server_pkdb.add_public_key("myclient1", 
            self.client1_pkdb.private_key.get_public_key())
        self.server_pkdb.add_public_key("myclient2", 
            self.client2_pkdb.private_key.get_public_key())
        self.client1_pkdb.add_public_key("localhost", 
            self.server_pkdb.private_key.get_public_key())
        
        self.log("starting server")
        self.server = ThreadedServer(rpyc.SlaveService, hostname = "localhost", 
            authenticator = PubkeyDBAuthenticator(self.server_pkdb))
        self.server.logger.quiet = True
        thread.start_new(self.server.start, ())
        time.sleep(1) # make sure the server has initialized, etc.
    
    def cleanup(self):
        self.server.close()
        time.sleep(1)
    
    def step_success(self):
        c = rpyc.classic.pubkey_connect("localhost", pkdb = self.client1_pkdb,
            port = self.server.port)
        self.log(c.modules.sys)
        self.log(len(c.modules.sys.path))
        c.close()

    def step_server_rejects_client(self):
        try:
            c = rpyc.classic.pubkey_connect("localhost", pkdb = self.client3_pkdb,
                port = self.server.port)
        except ProtocolAuthError, ex:
            self.log(ex)
        else:
            self.fail("server accepted us")

    def step_client_rejects_server(self):
        try:
            c = rpyc.classic.pubkey_connect("localhost", pkdb = self.client2_pkdb,
                port = self.server.port)
        except ProtocolAuthError, ex:
            self.log(ex)
        else:
            self.fail("client accepted server")

    def step_client_accepts_unknown_server(self):
        c = rpyc.classic.pubkey_connect("localhost", pkdb = self.client2_pkdb,
            port = self.server.port, allow_unknown_host = True)
        self.log(c.modules.sys)
        self.log(len(c.modules.sys.path))
        c.close()

if __name__ == "__main__":
    PkdbTest.run()





